import torch
import torchvision.datasets as dataset
import torchvision.transforms as transform
import torch.utils.data as data_utils

# data
# 加载MNIST训练数据集
train_data = dataset.MNIST(
    root="./mnist/", 
    train=True,  # 是否为训练集
    transform=transform.ToTensor(),  # 将图像转换为Tensor格式
    download=True,  # 是否下载数据集
)
# 加载MNIST测试数据集
test_data = dataset.MNIST(
    root="./mnist/", 
    train=False, 
    transform=transform.ToTensor(), 
    download=False,
)

# batchsize
# 创建训练数据加载器，设置数据集、批量大小为64以及是否打乱数据
train_loader = data_utils.DataLoader(
    dataset=train_data, 
    batch_size=64, 
    shuffle=True  # 避免因数据顺序而产生的偏差，提高泛化能力

) 
# 创建测试数据加载器，设置数据集、批量大小为64以及是否打乱数据
test_loader = data_utils.DataLoader(dataset=test_data, batch_size=64, shuffle=True)

# net
# 定义一个卷积神经网络类，继承自torch.nn.Module
class CNN(torch.nn.Module):
    def __init__(self):
        # 调用父类的构造函数
        super(CNN, self).__init__()
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(1, 32, kernel_size=5, padding=2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )

        self.fc = torch.nn.Linear(14 * 14 * 32, 10)

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size()[0], -1)
        out = self.fc(out)

        return out

cnn = CNN() # 初始化类

# loss
loss_func = torch.nn.CrossEntropyLoss() # 分类 采用交叉熵损失函数

# optimizer
optimizer = torch.optim.Adam(cnn.parameters(), lr=0.01)

# training
for epoch in range(10):
    for i, (images, labels) in enumerate(train_loader):
        outputs = cnn(images)
        loss = loss_func(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # 每个迭代打印一次训练信息（将 i 从 0-base 转为 1-base），格式化为：epoch, 当前迭代/总迭代数, loss
    print(f"epoch is {epoch+1}, ite is {i+1}/{len(train_data) // 64}, loss is {loss.item():.6f}")

    # eval/test
    loss_test = 0
    accuracy = 0
    for i, (images, labels) in enumerate(test_loader):
        outputs = cnn(images)
        # [batchsize]
        # outputs = batchsize * cls_num
        loss_test += loss_func(outputs, labels)
        _, pred = outputs.max(1)
        accuracy += (pred == labels).sum().item()

    accuracy = accuracy / len(test_data)
    loss_test = loss_test / (len(test_data) // 64)

    print("epoch is {}, accuracy is {}, test loss is {}".format(epoch+1, accuracy, loss_test))

torch.save(cnn, "model/mnist_model.pkl")
